import torch
import torch.nn as nn

def pair_loss(
    out, gt_map, *, 
    gamma=32, margin = 0.25, neg_rate=-1, neg_num=-1
):
    """
    sample pair loss

    Param:
    -----
    neg_rate :  rate between negative sample number and the 
                postives (invalid when neg_num>0)
                -1, all
                1, same with the positives
    neg_num :   number of negative sample
                -1, no usage
    """


    gt_map_neg = (1-gt_map).type(dtype=torch.bool)
    gt_mapx = gt_map.type(dtype=torch.bool)
    sp = out[gt_mapx]
    sn = out[gt_map_neg]

    sp = sp.view(out.size()[0], -1)
    sn = sn.view(out.size()[0], -1)

    pos_num = sp.size(1)
    if neg_num == -1:
        if neg_rate == -1:
            neg_num = sn.size(1)
        else:
            neg_num = min(int(pos_num*neg_rate), sn.size(1))

    print(neg_num)

    sorted_sn, _ = torch.sort(sn, -1, descending=True)

    sn = sorted_sn[:, :neg_num]

    ap = torch.clamp_min(-sp.detach() + 1 + margin, min=0.)
    an = torch.clamp_min(sn.detach() + margin, min=0.)

    delta_p = 1-margin
    delta_n = margin

    logit_p = -ap * (sp - delta_p) * gamma
    logit_n = an * (sn - delta_n) * gamma

    soft_plus = nn.Softplus()
    loss_circle = soft_plus(torch.logsumexp(logit_n, dim=1) + torch.logsumexp(logit_p, dim=1))

    loss_regu = loss_circle.mean()

    return loss_regu

def test_(x, y, done=1):

    print(done)
